import os
import Utility.Stl as Stl
import Utility.tetrahedronTXT as Txt
import matplotlib.pyplot as plt

def show_grain_from_stl_and_txt(grain_index : int, data_path : str):
    # 确保data_path路径存在
    os.makedirs(data_path, exist_ok=True)

    # 提取信息STL文件名词信息（例如0#N6L3215BS.stl）
    extraction = Stl.extract_info_from_stl_filename(os.listdir(data_path), grain_index)

    # 读取STL文件，获取初始燃面点云
    bs_points_cloud = Stl.read_stl_by_given_point_num(f"{data_path}/{extraction['stl_full_name']}")
    
    # 读取TXT文件，获取空腔区域点云
    cr_points_cloud = Txt.read_txt_by_given_point_num(f"{data_path}/{extraction['cr_txt_full_name']}")

    # 读取TXT文件，获取固体区域点云
    sr_points_cloud = Txt.read_txt_by_given_point_num(f"{data_path}/{extraction['sr_txt_full_name']}")

    # 创建3D图形
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制点云
    ax.scatter(bs_points_cloud[:, 0], bs_points_cloud[:, 1], bs_points_cloud[:, 2], c='r', marker='.', label='BS')
    ax.scatter(cr_points_cloud[:, 0], cr_points_cloud[:, 1], cr_points_cloud[:, 2], c='g', marker='.', label='CR')
    ax.scatter(sr_points_cloud[:, 0], sr_points_cloud[:, 1], sr_points_cloud[:, 2], c='b', marker='.', label='SR')

    # 设置坐标轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # 设置坐标轴比例尺相同
    ax.axis('equal')

    # 添加图例
    ax.legend()

    # 显示图形
    plt.show()

if __name__ == '__main__':
    show_grain_from_stl_and_txt(240, "/media/hpc/Backup_lwt/3DModel/G478")